// K6 Load Testing Script for ArXiv Subscription Platform
// Run with: k6 run --vus 100 --duration 5m load_test.js

import http from 'k6/http';
import { check, sleep, group } from 'k6';
import { Rate, Trend, Counter } from 'k6/metrics';

// Custom metrics
let errorRate = new Rate('errors');
let responseTime = new Trend('response_time');
let requestCounter = new Counter('requests');

// Test configuration
export let options = {
  stages: [
    { duration: '2m', target: 10 },   // Ramp up to 10 users
    { duration: '5m', target: 50 },   // Stay at 50 users
    { duration: '10m', target: 100 }, // Ramp up to 100 users
    { duration: '15m', target: 100 }, // Stay at 100 users
    { duration: '5m', target: 200 },  // Peak load test
    { duration: '3m', target: 200 },  // Stay at peak
    { duration: '5m', target: 0 },    // Ramp down
  ],
  thresholds: {
    http_req_duration: ['p(95)<500'], // 95% of requests should be under 500ms
    http_req_failed: ['rate<0.01'],   // Error rate should be less than 1%
    errors: ['rate<0.01'],
  },
};

// Test data
const BASE_URL = __ENV.BASE_URL || 'https://localhost';
const API_BASE = `${BASE_URL}/api`;

// Test user credentials
const testUsers = [
  { email: 'test1@example.com', password: 'TestPass123!' },
  { email: 'test2@example.com', password: 'TestPass123!' },
  { email: 'test3@example.com', password: 'TestPass123!' },
  { email: 'test4@example.com', password: 'TestPass123!' },
  { email: 'test5@example.com', password: 'TestPass123!' },
];

// Sample search terms
const searchTerms = [
  'machine learning',
  'quantum computing',
  'neural networks',
  'artificial intelligence',
  'computer vision',
  'natural language processing',
  'deep learning',
  'reinforcement learning',
  'data science',
  'statistics'
];

// Sample ArXiv categories
const categories = [
  'cs.AI', 'cs.LG', 'cs.CV', 'cs.CL', 'cs.NE',
  'stat.ML', 'math.ST', 'physics.data-an',
  'q-bio.QM', 'econ.EM'
];

// Helper function to get random element from array
function getRandomElement(arr) {
  return arr[Math.floor(Math.random() * arr.length)];
}

// Helper function to authenticate user
function authenticateUser() {
  const user = getRandomElement(testUsers);
  
  let response = http.post(`${API_BASE}/auth/login`, JSON.stringify({
    email: user.email,
    password: user.password
  }), {
    headers: {
      'Content-Type': 'application/json',
    },
  });
  
  requestCounter.add(1);
  responseTime.add(response.timings.duration);
  
  if (response.status === 200 && response.json('access_token')) {
    return response.json('access_token');
  }
  
  errorRate.add(1);
  return null;
}

// Helper function to create authenticated headers
function getAuthHeaders(token) {
  return {
    'Content-Type': 'application/json',
    'Authorization': `Bearer ${token}`,
  };
}

// Main test function
export default function() {
  const token = authenticateUser();
  if (!token) {
    console.error('Failed to authenticate');
    return;
  }
  
  const headers = getAuthHeaders(token);
  
  group('Homepage Load', () => {
    let response = http.get(BASE_URL);
    requestCounter.add(1);
    responseTime.add(response.timings.duration);
    
    check(response, {
      'homepage loads': (r) => r.status === 200,
      'homepage has title': (r) => r.body.includes('ArXiv'),
    }) || errorRate.add(1);
    
    sleep(1);
  });
  
  group('API Health Check', () => {
    let response = http.get(`${API_BASE}/health`);
    requestCounter.add(1);
    responseTime.add(response.timings.duration);
    
    check(response, {
      'health check returns 200': (r) => r.status === 200,
      'health check has status': (r) => r.json('status') === 'healthy',
    }) || errorRate.add(1);
  });
  
  group('User Dashboard', () => {
    let response = http.get(`${API_BASE}/dashboard`, { headers });
    requestCounter.add(1);
    responseTime.add(response.timings.duration);
    
    check(response, {
      'dashboard loads': (r) => r.status === 200,
      'dashboard has recommendations': (r) => r.json('recommendations'),
    }) || errorRate.add(1);
    
    sleep(1);
  });
  
  group('Paper Search', () => {
    const searchTerm = getRandomElement(searchTerms);
    
    let response = http.get(`${API_BASE}/papers/search?q=${encodeURIComponent(searchTerm)}&limit=20`, { headers });
    requestCounter.add(1);
    responseTime.add(response.timings.duration);
    
    check(response, {
      'search returns 200': (r) => r.status === 200,
      'search returns results': (r) => r.json('papers') && r.json('papers').length > 0,
      'search has pagination': (r) => r.json('total_count') !== undefined,
    }) || errorRate.add(1);
    
    sleep(2);
  });
  
  group('Category Browse', () => {
    const category = getRandomElement(categories);
    
    let response = http.get(`${API_BASE}/papers?category=${category}&limit=20`, { headers });
    requestCounter.add(1);
    responseTime.add(response.timings.duration);
    
    check(response, {
      'category browse returns 200': (r) => r.status === 200,
      'category has papers': (r) => r.json('papers') && r.json('papers').length > 0,
    }) || errorRate.add(1);
    
    sleep(1);
  });
  
  group('Paper Details', () => {
    // First get a list of papers
    let listResponse = http.get(`${API_BASE}/papers?limit=5`, { headers });
    requestCounter.add(1);
    responseTime.add(listResponse.timings.duration);
    
    if (listResponse.status === 200 && listResponse.json('papers').length > 0) {
      const paperId = listResponse.json('papers')[0].id;
      
      let detailResponse = http.get(`${API_BASE}/papers/${paperId}`, { headers });
      requestCounter.add(1);
      responseTime.add(detailResponse.timings.duration);
      
      check(detailResponse, {
        'paper detail returns 200': (r) => r.status === 200,
        'paper has title': (r) => r.json('title'),
        'paper has authors': (r) => r.json('authors') && r.json('authors').length > 0,
      }) || errorRate.add(1);
    }
    
    sleep(1);
  });
  
  group('User Interactions', () => {
    // Get a paper to interact with
    let listResponse = http.get(`${API_BASE}/papers?limit=1`, { headers });
    requestCounter.add(1);
    responseTime.add(listResponse.timings.duration);
    
    if (listResponse.status === 200 && listResponse.json('papers').length > 0) {
      const paperId = listResponse.json('papers')[0].id;
      
      // Bookmark paper
      let bookmarkResponse = http.post(`${API_BASE}/bookmarks`, JSON.stringify({
        paper_id: paperId
      }), { headers });
      requestCounter.add(1);
      responseTime.add(bookmarkResponse.timings.duration);
      
      check(bookmarkResponse, {
        'bookmark creation succeeds': (r) => r.status === 201 || r.status === 200,
      }) || errorRate.add(1);
      
      // Rate paper
      let ratingResponse = http.post(`${API_BASE}/ratings`, JSON.stringify({
        paper_id: paperId,
        rating: Math.floor(Math.random() * 5) + 1
      }), { headers });
      requestCounter.add(1);
      responseTime.add(ratingResponse.timings.duration);
      
      check(ratingResponse, {
        'rating creation succeeds': (r) => r.status === 201 || r.status === 200,
      }) || errorRate.add(1);
    }
    
    sleep(1);
  });
  
  group('Subscription Management', () => {
    // Create subscription
    const category = getRandomElement(categories);
    
    let createResponse = http.post(`${API_BASE}/subscriptions`, JSON.stringify({
      subscription_type: 'category',
      subscription_target: category,
      subscription_name: `${category} Papers`
    }), { headers });
    requestCounter.add(1);
    responseTime.add(createResponse.timings.duration);
    
    check(createResponse, {
      'subscription creation succeeds': (r) => r.status === 201 || r.status === 200,
    }) || errorRate.add(1);
    
    // Get user subscriptions
    let listResponse = http.get(`${API_BASE}/subscriptions`, { headers });
    requestCounter.add(1);
    responseTime.add(listResponse.timings.duration);
    
    check(listResponse, {
      'subscriptions list returns 200': (r) => r.status === 200,
      'subscriptions list is array': (r) => Array.isArray(r.json()),
    }) || errorRate.add(1);
    
    sleep(1);
  });
  
  group('Email Preferences', () => {
    let response = http.get(`${API_BASE}/email-preferences`, { headers });
    requestCounter.add(1);
    responseTime.add(response.timings.duration);
    
    check(response, {
      'email preferences returns 200': (r) => r.status === 200,
      'has digest frequency': (r) => r.json('digest_frequency'),
    }) || errorRate.add(1);
    
    // Update preferences
    let updateResponse = http.put(`${API_BASE}/email-preferences`, JSON.stringify({
      digest_frequency: 'weekly',
      instant_notifications: Math.random() > 0.5
    }), { headers });
    requestCounter.add(1);
    responseTime.add(updateResponse.timings.duration);
    
    check(updateResponse, {
      'preference update succeeds': (r) => r.status === 200,
    }) || errorRate.add(1);
    
    sleep(1);
  });
  
  // Random think time between 1-5 seconds
  sleep(Math.random() * 4 + 1);
}

// Setup function (runs once per VU at the beginning)
export function setup() {
  console.log('Starting load test setup...');
  
  // Check if the service is available
  let response = http.get(`${BASE_URL}/health`);
  if (response.status !== 200) {
    throw new Error(`Service not available. Status: ${response.status}`);
  }
  
  console.log('Service is available, starting load test...');
  return { baseUrl: BASE_URL };
}

// Teardown function (runs once after all VUs finish)
export function teardown(data) {
  console.log('Load test completed');
  console.log(`Base URL: ${data.baseUrl}`);
}

// Handle summary to output custom metrics
export function handleSummary(data) {
  return {
    'stdout': textSummary(data, { indent: ' ', enableColors: true }),
    'load_test_results.json': JSON.stringify(data, null, 2),
  };
}

function textSummary(data, options = {}) {
  const indent = options.indent || '';
  const enableColors = options.enableColors || false;
  
  let summary = `
${indent}Test Summary
${indent}============
${indent}Total Requests: ${data.metrics.requests.values.count}
${indent}Failed Requests: ${data.metrics.http_req_failed.values.rate * 100}%
${indent}Average Response Time: ${data.metrics.http_req_duration.values.avg}ms
${indent}95th Percentile: ${data.metrics.http_req_duration.values['p(95)']}ms
${indent}99th Percentile: ${data.metrics.http_req_duration.values['p(99)']}ms
${indent}Max Response Time: ${data.metrics.http_req_duration.values.max}ms
${indent}Requests per Second: ${data.metrics.http_reqs.values.rate}
${indent}Data Received: ${(data.metrics.data_received.values.count / 1024 / 1024).toFixed(2)}MB
${indent}Data Sent: ${(data.metrics.data_sent.values.count / 1024 / 1024).toFixed(2)}MB

${indent}Thresholds
${indent}==========
`;

  for (const [name, threshold] of Object.entries(data.thresholds || {})) {
    const passed = threshold.ok ? '✓' : '✗';
    summary += `${indent}${passed} ${name}\n`;
  }

  return summary;
}